#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
双色球一键训练脚本
自动完成数据采集、基础LSTM模型训练和增强版LSTM-CRF模型训练
"""

import os
import sys
import subprocess
import time
from datetime import datetime

# 添加项目根目录到路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = current_dir
sys.path.append(project_root)

def print_header(title):
    """打印标题"""
    print("=" * 60)
    print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] {title}")
    print("=" * 60)

def run_script(script_path, description):
    """运行脚本并检查结果"""
    print(f"\n[PROCESS] 开始{description}...")
    print(f"[SCRIPT] 执行: python {script_path}")
    
    try:
        # 使用subprocess运行脚本
        result = subprocess.run(
            [sys.executable, script_path], 
            cwd=project_root,
            capture_output=True, 
            text=True,
            timeout=300  # 5分钟超时
        )
        
        # 打印输出
        if result.stdout:
            print(result.stdout)
        
        if result.stderr:
            print(result.stderr)
        
        # 检查返回码
        if result.returncode == 0:
            print(f"[SUCCESS] {description}完成！")
            return True
        else:
            print(f"[ERROR] {description}失败！")
            print(f"返回码: {result.returncode}")
            return False
            
    except subprocess.TimeoutExpired:
        print(f"[ERROR] {description}超时！")
        return False
    except Exception as e:
        print(f"[ERROR] 运行{description}时出错: {e}")
        return False

def check_file_exists(file_path, description):
    """检查文件是否存在"""
    if os.path.exists(file_path):
        size = os.path.getsize(file_path)
        print(f"[CHECK] {description}存在 (大小: {size} 字节)")
        return True
    else:
        print(f"[CHECK] {description}不存在")
        return False

def main():
    """主函数"""
    print_header("双色球一键训练程序")
    
    # 双色球相关路径
    ssq_dir = os.path.join(project_root, 'scripts', 'ssq')
    history_file = os.path.join(ssq_dir, 'ssq_history.csv')
    lstm_model_file = os.path.join(ssq_dir, 'ssq_lstm_model.pth')
    enhanced_model_file = os.path.join(ssq_dir, 'enhanced_lstm_crf_model.pth')
    
    # 1. 检查历史数据
    print("\n[STEP 1/4] 检查双色球历史数据...")
    if not check_file_exists(history_file, "双色球历史数据文件"):
        # 如果历史数据不存在，先获取数据
        fetch_script = os.path.join(ssq_dir, 'fetch_ssq_data.py')
        if os.path.exists(fetch_script):
            if not run_script(fetch_script, "获取双色球历史数据"):
                print("[ERROR] 无法获取双色球历史数据，程序终止！")
                return False
        else:
            print(f"[ERROR] 数据获取脚本不存在: {fetch_script}")
            return False
    
    # 再次检查历史数据
    if not check_file_exists(history_file, "双色球历史数据文件"):
        print("[ERROR] 双色球历史数据文件不存在，程序终止！")
        return False
    
    # 2. 训练基础LSTM模型
    print("\n[STEP 2/4] 训练基础LSTM模型...")
    if not check_file_exists(lstm_model_file, "基础LSTM模型文件"):
        train_script = os.path.join(ssq_dir, 'train_ssq_model.py')
        if os.path.exists(train_script):
            if not run_script(train_script, "训练基础LSTM模型"):
                print("[ERROR] 基础LSTM模型训练失败，程序终止！")
                return False
        else:
            print(f"[ERROR] 基础LSTM模型训练脚本不存在: {train_script}")
            return False
    
    # 检查基础LSTM模型是否生成
    if not check_file_exists(lstm_model_file, "基础LSTM模型文件"):
        print("[ERROR] 基础LSTM模型文件未生成，程序终止！")
        return False
    
    # 3. 训练增强版LSTM-CRF模型
    print("\n[STEP 3/4] 训练增强版LSTM-CRF模型...")
    if not check_file_exists(enhanced_model_file, "增强版LSTM-CRF模型文件"):
        enhanced_train_script = os.path.join(ssq_dir, 'train_enhanced_lstm_crf.py')
        if os.path.exists(enhanced_train_script):
            if not run_script(enhanced_train_script, "训练增强版LSTM-CRF模型"):
                print("[ERROR] 增强版LSTM-CRF模型训练失败，程序终止！")
                return False
        else:
            print(f"[ERROR] 增强版LSTM-CRF模型训练脚本不存在: {enhanced_train_script}")
            return False
    
    # 4. 最终检查
    print("\n[STEP 4/4] 最终检查...")
    success = True
    
    if not check_file_exists(enhanced_model_file, "增强版LSTM-CRF模型文件"):
        print("[WARN] 增强版LSTM-CRF模型文件未生成")
        success = False
    
    if success:
        print_header("双色球模型训练全部完成！")
        print("现在可以使用双色球💎 LSTM-CRF 经典模式增强版进行预测了。")
        return True
    else:
        print_header("双色球模型训练未完全成功！")
        return False

if __name__ == "__main__":
    try:
        success = main()
        if success:
            print("\n[COMPLETE] 所有步骤执行成功！")
            sys.exit(0)
        else:
            print("\n[FAIL] 部分步骤执行失败！")
            sys.exit(1)
    except KeyboardInterrupt:
        print("\n[INTERRUPT] 用户中断程序执行")
        sys.exit(1)
    except Exception as e:
        print(f"\n[ERROR] 程序执行出错: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)